#include <cuda_runtime.h>
#include <device_launch_parameters.h>
#include <stdio.h>
#include <stdlib.h>
#include <string.h>
#include <math.h>
#include <time.h>
#include <sys/time.h>

const int TILE_WIDTH = 16;

double get_cpu_time_ms() {
	struct timespec ts;
	clock_gettime(CLOCK_MONOTONIC, &ts);
	return (ts.tv_sec * 1000.0) + (ts.tv_nsec / 1000000.0);
}

void TransposeCPU(const double* A, double* B, int rows, int cols) {
	for (int r = 0; r < rows; ++r) {
		for (int c = 0; c < cols; ++c) {
			B[c * rows + r] = A[r * cols + c];
		}
	}
}

__global__ void TransposeKernel(const double* __restrict__ A,
								double* __restrict__ B,
								int rows, int cols, bool useShared) {
	if (!useShared) {
		int c = blockIdx.x * blockDim.x + threadIdx.x; 
		int r = blockIdx.y * blockDim.y + threadIdx.y; 
		if (r < rows && c < cols) {
			B[c * rows + r] = A[r * cols + c];
		}
		return;
	}

	__shared__ double tile[TILE_WIDTH][TILE_WIDTH];
    int bx = blockIdx.x;
    int by = blockIdx.y;
    int tx = threadIdx.x;
    int ty = threadIdx.y;
	int x = bx * TILE_WIDTH + tx;
	int y = by * TILE_WIDTH + ty;

	if (y < rows && x < cols) {
		tile[ty][tx] = A[y * cols + x];
	} else {
		tile[ty][tx] = 0.0;
	}
	__syncthreads();

	int xT = by * TILE_WIDTH + tx; 
	int yT = bx * TILE_WIDTH + ty;

	if (yT < cols && xT < rows) {
		B[yT * rows + xT] = tile[threadIdx.x][threadIdx.y];
	}
}

static void verify(const double* ref, const double* out, size_t len) {
	size_t err_cnt = 0;
	const double eps = 1e-9;
	for (size_t i = 0; i < len; ++i) {
		double e = fabs(ref[i] - out[i]);
		if (e > eps) err_cnt++;
	}
	if (err_cnt) {
		printf("Verification failed! Number of errors: %zu\n", err_cnt);
	} else {
		printf("Verification passed!\n");
	}
}

int main() {
	const char* inputPath = "input2.txt";
	const char* outputPath = "output2.txt";

	int m = 0, n = 0;
	FILE* fin = fopen(inputPath, "r");
    fscanf(fin, "%d,%d", &m, &n);
    fclose(fin);

	size_t sizeA = (size_t)m * (size_t)n;
	size_t sizeB = (size_t)n * (size_t)m;
	double* h_A = (double*)malloc(sizeA * sizeof(double));
	double* h_P_cpu = (double*)malloc(sizeB * sizeof(double));
	double* h_P_naive = (double*)malloc(sizeB * sizeof(double));
	double* h_P_shared = (double*)malloc(sizeB * sizeof(double));

	for (size_t i = 0; i < sizeA; ++i) {
		h_A[i] = (double)rand() / (double)RAND_MAX;
	}

	double cpu_start = get_cpu_time_ms();
	TransposeCPU(h_A, h_P_cpu, m, n);
	double cpu_stop = get_cpu_time_ms();
	double cpu_time = cpu_stop - cpu_start;
	printf("CPU Time: %.2f ms\n", cpu_time);

	double *d_A = nullptr, *d_P_naive = nullptr, *d_P_shared = nullptr;
	cudaMalloc((void**)&d_A, sizeA * sizeof(double));
	cudaMalloc((void**)&d_P_naive, sizeB * sizeof(double));
    cudaMalloc((void**)&d_P_shared, sizeB * sizeof(double));
	cudaMemcpy(d_A, h_A, sizeA * sizeof(double), cudaMemcpyHostToDevice);

	dim3 block(TILE_WIDTH, TILE_WIDTH);
	dim3 grid((n + TILE_WIDTH - 1) / TILE_WIDTH, (m + TILE_WIDTH - 1) / TILE_WIDTH);

	cudaEvent_t startEvent, stopEvent;
	cudaEventCreate(&startEvent);
	cudaEventCreate(&stopEvent);

	cudaMemset(d_P_naive, 0, sizeB * sizeof(double));
	cudaEventRecord(startEvent);
	TransposeKernel<<<grid, block>>>(d_A, d_P_naive, m, n, false);
	cudaEventRecord(stopEvent);
	cudaEventSynchronize(stopEvent);
	float elapsedNaive = 0.0f;
	cudaEventElapsedTime(&elapsedNaive, startEvent, stopEvent);
	double gpu_time_naive = (double)elapsedNaive;
	printf("GPU Time (no shared): %.2f ms\n", gpu_time_naive);
	cudaMemcpy(h_P_naive, d_P_naive, sizeB * sizeof(double), cudaMemcpyDeviceToHost);
	verify(h_P_cpu, h_P_naive, sizeB);

	cudaMemset(d_P_shared, 0, sizeB * sizeof(double));
	cudaEventRecord(startEvent);
	TransposeKernel<<<grid, block>>>(d_A, d_P_shared, m, n, true);
	cudaEventRecord(stopEvent);
	cudaEventSynchronize(stopEvent);
	float elapsedShared = 0.0f;
	cudaEventElapsedTime(&elapsedShared, startEvent, stopEvent);
	double gpu_time_shared = (double)elapsedShared;
	printf("GPU Time (shared): %.2f ms\n", gpu_time_shared);
	cudaMemcpy(h_P_shared, d_P_shared, sizeB * sizeof(double), cudaMemcpyDeviceToHost);
	verify(h_P_cpu, h_P_shared, sizeB);

	FILE* fout = fopen(outputPath, "w");
	if (fout) {
		fprintf(fout, "%.2f,%.2f,%.2f", cpu_time, gpu_time_naive, gpu_time_shared);
		fclose(fout);
	} else {
		fprintf(stderr, "Failed to open %s for writing\n", outputPath);
	}

	cudaEventDestroy(startEvent);
	cudaEventDestroy(stopEvent);
	cudaFree(d_A);
	cudaFree(d_P_naive);
	cudaFree(d_P_shared);
	free(h_A);
	free(h_P_cpu);
	free(h_P_naive);
	free(h_P_shared);
	return 0;
}

